import os
import librosa
import math
import json

DATASET_PATH = "D:\\Audio\\datasets\\genres"
JSON_PATH = "D:\\Audio\\datasets\\genres\\data_10.json"

SAMPLE_RATE = 22050
DURATION = 30  # seconds
SAMPLES_PER_TRACK = SAMPLE_RATE * DURATION

def save_mfcc(dataset_path, json_path, 
              n_mfcc=13, n_fft=2048, hop_length=512, num_segments=5):
    """
    num_segments: 每首歌分成的片段
    """
    # dictionary to store data
    # data = {
    #     "mapping": ["classical", "blues"],  # genres
    #     "mfcc": [[...], [...], [...]],
    #     "labels": [0, 0, 1]
    # }
    data = {
        "mapping": [],
        "mfcc": [],
        "labels": []
    }
    
    num_samples_per_segment = int(SAMPLES_PER_TRACK / num_segments)
    # 
    excepted_num_mfcc_vectors_per_segment = math.ceil(num_samples_per_segment / hop_length)
    
    # loop through all the genres
    for i, (dirpath, dirnames, filenames) in enumerate(os.walk(dataset_path)):
        
        # ensure that we're not at the root level
        if dirpath is not dataset_path:
            
            # save the semantic label
            dirpath_components = dirpath.split("\\")  # genre/blues => ["genre", "blues"]
            semantic_label = dirpath_components[-1]
            data["mapping"].append(semantic_label)
            print("\nProcessing {}".format(semantic_label))
            
            # process files for a specific genre
            for f in filenames:
                
                # load the audio file
                file_path = os.path.join(dirpath, f)
                signal, sr = librosa.load(file_path, sr=SAMPLE_RATE)
                
                # process segments extracting MFCC and storing data
                for s in range(num_segments):
                    start_sample = num_samples_per_segment * s
                    finish_sample = start_sample + num_samples_per_segment
                    
                    mfcc = librosa.feature.mfcc(signal[start_sample:finish_sample],
                                                sr=sr,
                                                n_fft=n_fft,
                                                n_mfcc=n_mfcc,
                                                hop_length=hop_length)
                    mfcc = mfcc.T 
                    
                    # store mfcc for segment if it has the expected(vector) length
                    if len(mfcc) == excepted_num_mfcc_vectors_per_segment:
                        data["mfcc"].append(mfcc.tolist())
                        data["labels"].append(i - 1)  # 忽略根目录
                        print("{}, segment:{}".format(file_path, s + 1))
                        
    with open(json_path, "w") as jp:
        json.dump(data, jp, indent=4)
        

if __name__ == "__main__":
    save_mfcc(DATASET_PATH, JSON_PATH, num_segments=10)
    